#!/usr/bin/env python3
# coding: utf-8
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import random
import torch
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig

@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul_swiglu_quant")
class AscendGroupedMatmulSwigluQuant(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        x = case_config.inputs[0]
        weight = case_config.inputs[1]
        weightScale = case_config.inputs[4]
        xScale = case_config.inputs[5]
        groupList = case_config.inputs[6]

        M = x.shape[0]

        #K需要是16，N需要是32
        weight.shape[3],weight.shape[4] = 16, 32
        E,K,N= weight.shape[0], weight.shape[2] * weight.shape[3], weight.shape[1] * weight.shape[4]
        
        #K轴不能大于等于65536
        if K >= 65536:
            weight.shape[2] = 4095
            weight.shape[3] = 16
        x.shape[1] = weight.shape[2] * weight.shape[3]
        #N轴不能大于8192
        if N > 10240:
            weight.shape[1],weight.shape[4] = 10240 // 32, 32
            N = 10240
        weightScale.shape[0] = E
        weightScale.shape[1] = N
        xScale.shape[0] = M
        # print(case_config)
        return case_config